import numpy as np
from PyEMD import EMD, Visualisation
import matplotlib.pyplot as plt
# 构建信号
# 时间t: 为0到1s,采样频率为100Hz，S为合成信号
from pandas import read_csv, DataFrame
from sklearn.preprocessing import MinMaxScaler
from statsmodels.graphics.tsaplots import plot_pacf
from tensorflow.python.keras.utils.np_utils import to_categorical
from keras.layers import LSTM, RepeatVector, Dense, \
    Activation, Add, Reshape, Input, Lambda, Multiply, Concatenate, Dot
from util import series_to_supervised, TIME_STEP, draw_data
import keras
from sklearn.metrics import mean_squared_error, mean_absolute_error
from tensorflow.python.keras.callbacks import EarlyStopping
def getFlowData():
    # 切换19年和20年数据只需要改下面两行
    multi_dataset = read_csv('./data/2019allday.csv', header=0, index_col=0)
    dataset = DataFrame()
    # 取in_card_flow(流出机场客流)、实际降落载客数 arr_ALDT_passenger、时段、天气作为参数，预测in_card_flow
    dataset['in_flow'] = multi_dataset['in_flow']
    # 将NAN替换为0
    dataset.fillna(0, inplace=True)
    values = dataset.values.flatten()
    return values

def getArr_ALDT():
    # 切换19年和20年数据只需要改下面两行
    multi_dataset = read_csv('./data/2019allday.csv', header=0, index_col=0)
    dataset = DataFrame()
    dataset['arr_ALDT_passenger'] = multi_dataset['arr_ALDT_passenger']
    # 将NAN替换为0
    dataset.fillna(0, inplace=True)
    values = dataset.values.flatten()
    return values
# PARAMETER_NUM: 分解出的本征函数个数+残波个数
def getModel(feature_num):
    inputs = Input(shape=(TIME_STEP, feature_num))  # 输入时间序列数据
    lstm = keras.layers.Bidirectional(LSTM(20, return_sequences=True))(inputs)
    lstm = keras.layers.Bidirectional(LSTM(20, return_sequences=True))(lstm)
    lstm = keras.layers.Bidirectional(LSTM(20, return_sequences=True))(lstm)
    lstm = keras.layers.Bidirectional(LSTM(20, return_sequences=True))(lstm)
    lstm = keras.layers.Bidirectional(LSTM(20, return_sequences=True))(lstm)
    output = Dense(1)(lstm)
    output = Reshape((-1,))(output)
    output = Dense(1)(output)  # 最后预测，代码修改点

    model = keras.models.Model(inputs=inputs, outputs=output)
    return model

if __name__ == '__main__':
    # 获取轨道客流数据
    flowData = getFlowData()
    day_num = 26
    # 提取imfs和残波
    emd = EMD()
    emd.emd(flowData)
    imfs, res = emd.get_imfs_and_residue()

    # 获取航班实际降落载客数据
    arr_ALDT = getArr_ALDT()
    emd = EMD()
    emd.emd(arr_ALDT)
    imfs_arr_ALDT, res_arr_ALDT = emd.get_imfs_and_residue()

    # 把所有的本征函数当做输入特征(残波数值基本为0，可以舍去)
    values = np.vstack((flowData, imfs))
    # values = np.array(imfs)
    input_num = values.shape[0]

    values = values.T  # 转置
    reframed = series_to_supervised(values, TIME_STEP, 1)  # 6步预测下一步
    # 丢弃我们并不想预测的列
    reframed.drop(reframed.columns[[-1, -2, -3, -4, -5, -6, -7, -8, -9, -10, -11]], axis=1, inplace=True)

    # 分割为训练集和测试集
    values = reframed.values
    n_train_time_slice = (day_num - 1) * 100
    inputs_train = values[:n_train_time_slice, :]
    inputs_test = values[n_train_time_slice:, :]

    # 重塑成3D形状 [样例, 时间步, 特征]
    inputs_train = inputs_train.reshape((inputs_train.shape[0], TIME_STEP, int(inputs_train.shape[1] / TIME_STEP)))
    inputs_test = inputs_test.reshape((inputs_test.shape[0], TIME_STEP, int(inputs_test.shape[1] / TIME_STEP)))


    # 构建输出数据
    output = flowData.tolist()[TIME_STEP:]
    output_train = output[:n_train_time_slice]
    output_test = output[n_train_time_slice:]

    # model.fit函数中传入的输入输出，必须是ndarray格式
    output_train = np.array(output_train)
    output_test = np.array(output_test)

    model = getModel(input_num)
    model.compile(loss='mse', optimizer='adam')
    # 拟合神经网络模型
    early_stopping = EarlyStopping(monitor='val_loss', patience=20, verbose=2)


    mean_loss = 0
    n = 10
    for i in range(n):
        model.fit(inputs_train, output_train, epochs=1000, batch_size=128, validation_data=(inputs_test, output_test),
                  callbacks=[early_stopping], verbose=2,
                  shuffle=False)
        yhat = model.predict(inputs_test)
        mse = mean_squared_error(output_test, yhat)
        mean_loss = mean_loss + mse
    mean_loss = mean_loss / n
    print('EMD+LSTM mean_MSE: %.3f' % mean_loss)
